import os
import pandas as pd
import numpy as np
import akshare as ak
from datetime import datetime

import qlib
from qlib.data.dataset import DatasetH
from qlib.contrib.data.handler import Alpha360
from qlib.contrib.model.pytorch_lstm import LSTM

PROVIDER_URI = r"h:\\2025github_project\\qlib-main\\akshare_cn_data"
INSTRUMENT = "SH600104"
START_TIME = "2020-01-01"
END_TIME = datetime.today().strftime("%Y-%m-%d")


def init_qlib():
    qlib.init(provider_uri=PROVIDER_URI, region="cn")


def build_dataset():
    handler_kwargs = {
        "start_time": START_TIME,
        "end_time": END_TIME,
        "fit_start_time": START_TIME,
        "fit_end_time": (pd.Timestamp(END_TIME) - pd.Timedelta(days=300)).strftime("%Y-%m-%d"),
        "instruments": [INSTRUMENT],
        "infer_processors": [
            {"class": "RobustZScoreNorm", "kwargs": {"fields_group": "feature", "clip_outlier": True}},
            {"class": "Fillna", "kwargs": {"fields_group": "feature"}},
        ],
        "learn_processors": [
            {"class": "DropnaLabel"},
            {"class": "CSRankNorm", "kwargs": {"fields_group": "label"}},
        ],
        # label aims at next-day return
        "label": ["Ref($close, -1) / $close - 1"],
    }
    handler = {
        "class": "Alpha360",
        "module_path": "qlib.contrib.data.handler",
        "kwargs": handler_kwargs,
    }
    segments = {
        "train": [START_TIME, (pd.Timestamp(END_TIME) - pd.Timedelta(days=365)).strftime("%Y-%m-%d")],
        "valid": [
            (pd.Timestamp(END_TIME) - pd.Timedelta(days=365)).strftime("%Y-%m-%d"),
            (pd.Timestamp(END_TIME) - pd.Timedelta(days=30)).strftime("%Y-%m-%d"),
        ],
        "test": [
            (pd.Timestamp(END_TIME) - pd.Timedelta(days=30)).strftime("%Y-%m-%d"),
            END_TIME,
        ],
    }
    return DatasetH(handler=handler, segments=segments)


def train_lstm(dataset):
    model = LSTM(
        d_feat=6,
        hidden_size=64,
        num_layers=2,
        dropout=0.0,
        n_epochs=50,
        lr=1e-3,
        batch_size=512,
        metric="loss",
        loss="mse",
        GPU=0,
    )
    model.fit(dataset)
    return model


def extract_dates_from_index(idx):
    if isinstance(idx, pd.MultiIndex):
        # try common names
        names = list(idx.names)
        if "datetime" in names:
            return pd.to_datetime(idx.get_level_values("datetime"))
        else:
            return pd.to_datetime(idx.get_level_values(-1))
    else:
        return pd.to_datetime(idx)


def compute_metrics(preds_ser):
    # get realized next-day returns from akshare
    df = ak.stock_zh_a_daily(symbol="sh600104", adjust="qfq").sort_values("date")
    df["date"] = pd.to_datetime(df["date"])
    df["ret_next"] = df["close"].shift(-1) / df["close"] - 1

    # align predictions by date
    pred_df = pd.DataFrame({
        "date": extract_dates_from_index(preds_ser.index),
        "pred": preds_ser.values,
    })
    merged = pred_df.merge(df[["date", "ret_next"]], on="date", how="inner").dropna()

    if len(merged) == 0:
        return None, None

    y_pred = merged["pred"].astype(float).values
    y_true = merged["ret_next"].astype(float).values

    # metrics
    rmse = float(np.sqrt(np.mean((y_pred - y_true) ** 2)))
    mae = float(np.mean(np.abs(y_pred - y_true)))
    # R^2 with protection for zero variance
    denom = np.sum((y_true - y_true.mean()) ** 2)
    r2 = float(1.0 - np.sum((y_pred - y_true) ** 2) / denom) if denom > 0 else np.nan
    # directional accuracy (hit ratio)
    sign_pred = np.sign(y_pred)
    sign_true = np.sign(y_true)
    mask = (sign_true != 0)  # exclude exact zeros to avoid ambiguity
    hit_ratio = float(np.mean(sign_pred[mask] == sign_true[mask])) if np.any(mask) else np.nan

    return {
        "samples": int(len(merged)),
        "start": merged["date"].min().strftime("%Y-%m-%d"),
        "end": merged["date"].max().strftime("%Y-%m-%d"),
        "rmse": rmse,
        "mae": mae,
        "r2": r2,
        "hit_ratio": hit_ratio,
    }, merged


def main():
    init_qlib()
    ds = build_dataset()
    model = train_lstm(ds)
    preds = model.predict(ds, segment="test")
    metrics, merged = compute_metrics(preds)

    if metrics is None:
        print("No overlapping dates between predictions and realized returns; cannot compute metrics.")
        return

    print("Evaluation on test segment:")
    print(f"Samples: {metrics['samples']} from {metrics['start']} to {metrics['end']}")
    print(f"RMSE: {metrics['rmse']:.6f}")
    print(f"MAE: {metrics['mae']:.6f}")
    print(f"R^2: {metrics['r2']:.6f}")
    print(f"Directional Accuracy (hit ratio): {metrics['hit_ratio']:.3f}")


if __name__ == "__main__":
    main()